# PowerShell Script: Create Conda Environment using Tsinghua Mirror
# 使用方法 (Usage):
# 1. List available environment files (列出可用环境文件):
#    .\create_conda_env.ps1 -List
# 2. Create environment from YAML file (使用YAML文件创建环境):
#    .\create_conda_env.ps1 -EnvFile "py39_compatible.yml" -EnvName "custom_name"
# 3. Create environment from requirements.txt (使用requirements.txt创建环境):
#    .\create_conda_env.ps1 -RequirementsFile "requirements.txt" -PythonVersion "3.9" -EnvName "custom_name"
# 4. Create without Tsinghua mirror (不使用清华镜像源创建):
#    .\create_conda_env.ps1 -RequirementsFile "requirements.txt" -PythonVersion "3.9" -NoTsinghua
# 5. Export environment dependencies (导出环境依赖):
#    .\create_conda_env.ps1 -Export -EnvName "environment_name"
# 6. Export all environments (导出所有环境依赖):
#    .\create_conda_env.ps1 -ExportAll

param (
    [string]$EnvFile,           # 环境文件名(YAML)
    [string]$EnvName,           # 自定义环境名称
    [switch]$List,              # 是否列出可用环境文件
    [switch]$NoTsinghua,        # 是否不使用清华镜像源
    [switch]$Export,            # 导出环境依赖
    [switch]$ExportAll,         # 导出所有环境依赖
    [string]$RequirementsFile,  # requirements.txt文件
    [string]$PythonVersion      # Python版本
)

# Tsinghua mirror URL (清华镜像源URL)
$TSINGHUA_MIRROR = "https://mirrors.tuna.tsinghua.edu.cn/anaconda"

# Environment directory (环境文件目录)
$ENV_DIR = "conda_environments"

# Configure SSL/TLS (配置SSL/TLS)
function Configure-SSL {
    # Force to use TLS 1.2 (强制使用TLS 1.2)
    [Net.ServicePointManager]::SecurityProtocol = [Net.SecurityProtocolType]::Tls12
}

# Configure conda channels (配置conda源)
function Configure-Channels {
    param (
        [bool]$UseTsinghua = $true
    )
    
    # Reset channels first (首先重置所有源)
    conda config --remove-key channels
    
    # Add conda-forge as default channel (添加conda-forge作为默认源)
    conda config --add channels conda-forge
    
    if ($UseTsinghua) {
        # Add Tsinghua mirrors with higher priority (添加优先级更高的清华镜像源)
        conda config --add channels ${TSINGHUA_MIRROR}/pkgs/main
        conda config --add channels ${TSINGHUA_MIRROR}/pkgs/free
        conda config --add channels ${TSINGHUA_MIRROR}/cloud/conda-forge
    }
    
    # Show channel URLs in output (在输出中显示通道URL)
    conda config --set show_channel_urls yes
}

# Export single environment (导出单个环境)
function Export-Environment {
    param (
        [string]$EnvName
    )
    
    if (-not (Test-Path $ENV_DIR)) {
        New-Item -Path $ENV_DIR -ItemType Directory -Force | Out-Null
    }
    
    $outputPath = Join-Path -Path $ENV_DIR -ChildPath "${EnvName}_requirements.txt"
    
    Write-Host "Exporting environment '$EnvName' to $outputPath..." -ForegroundColor Cyan
    
    try {
        # Export pip packages (导出pip包)
        $command = "conda activate $EnvName && pip freeze > `"$outputPath`""
        Write-Host "Executing: $command" -ForegroundColor Gray
        
        # We need to use cmd to execute the conda activate command
        $fullCommand = "cmd.exe /c `"conda activate $EnvName && pip freeze > `"$outputPath`"`""
        Invoke-Expression $fullCommand
        
        if (Test-Path $outputPath) {
            Write-Host "Environment '$EnvName' successfully exported to $outputPath" -ForegroundColor Green
            return $true
        } else {
            Write-Host "Failed to export environment '$EnvName'" -ForegroundColor Red
            return $false
        }
    }
    catch {
        Write-Host "Error exporting environment '$EnvName': $_" -ForegroundColor Red
        return $false
    }
}

# Export environment as YAML (导出环境为YAML文件)
function Export-EnvironmentYAML {
    param (
        [string]$EnvName
    )
    
    if (-not (Test-Path $ENV_DIR)) {
        New-Item -Path $ENV_DIR -ItemType Directory -Force | Out-Null
    }
    
    $outputPath = Join-Path -Path $ENV_DIR -ChildPath "${EnvName}_environment.yml"
    
    Write-Host "Exporting environment '$EnvName' to YAML at $outputPath..." -ForegroundColor Cyan
    
    try {
        # 使用 --no-builds 参数提高跨平台兼容性
        $command = "conda env export --no-builds -n $EnvName > `"$outputPath`""
        Write-Host "Executing: $command" -ForegroundColor Gray
        
        Invoke-Expression $command
        
        if (Test-Path $outputPath) {
            # 读取并修改YAML文件内容，去除平台特定信息
            $content = Get-Content -Path $outputPath -Raw
            
            # 移除前缀路径
            $content = $content -replace 'prefix: .*', ''
            
            # 重新写入文件
            Set-Content -Path $outputPath -Value $content
            
            Write-Host "Environment '$EnvName' successfully exported to YAML at $outputPath" -ForegroundColor Green
            return $true
        } else {
            Write-Host "Failed to export environment '$EnvName' to YAML" -ForegroundColor Red
            return $false
        }
    }
    catch {
        Write-Host "Error exporting environment '$EnvName' to YAML: $_" -ForegroundColor Red
        return $false
    }
}

# Export all environments (导出所有环境)
function Export-AllEnvironments {
    $environments = & conda env list | Where-Object { $_ -match "^\w" } | ForEach-Object {
        $envLine = $_ -split "\s+"
        if ($envLine.Count -ge 2) {
            return $envLine[0] # 环境名称
        }
    }
    
    $environments = $environments | Where-Object { $_ -ne "base" -and $_ }
    
    if ($environments.Count -eq 0) {
        Write-Host "No environments found to export" -ForegroundColor Yellow
        return
    }
    
    Write-Host "Found $($environments.Count) environments to export" -ForegroundColor Cyan
    
    foreach ($env in $environments) {
        Write-Host "Processing environment: $env" -ForegroundColor White
        Export-Environment -EnvName $env
        Export-EnvironmentYAML -EnvName $env
    }
    
    Write-Host "All environments exported to $ENV_DIR directory" -ForegroundColor Green
}

# List available environment files (列出可用的环境文件)
function List-AvailableEnvironments {
    if (-not (Test-Path $ENV_DIR)) {
        Write-Host "Error: Environment directory $ENV_DIR not found" -ForegroundColor Red
        return @()
    }
    
    $env_files = Get-ChildItem -Path $ENV_DIR -Filter "*.yml" | Select-Object -ExpandProperty Name
    return $env_files
}

# List conda environments (列出所有conda环境)
function List-CondaEnvironments {
    $environments = & conda env list | Where-Object { $_ -match "^\w" } | ForEach-Object {
        $envLine = $_ -split "\s+"
        if ($envLine.Count -ge 2) {
            [PSCustomObject]@{
                Name = $envLine[0]
                Path = $envLine[1]
            }
        }
    }
    
    return $environments
}

# Create environment from requirements.txt or YAML (创建环境)
function Create-Environment {
    param (
        [string]$FilePath,           # 文件路径
        [string]$CustomEnvName,      # 自定义环境名称
        [string]$PythonVersion,      # Python版本
        [bool]$UseTsinghua = $true,  # 是否使用清华镜像源
        [bool]$IsYaml = $false       # 是否是YAML文件
    )
    
    if (-not (Test-Path $FilePath)) {
        Write-Host "Error: File $FilePath not found" -ForegroundColor Red
        return $false
    }
    
    # 如果没有指定环境名称，提示用户输入
    if (-not $CustomEnvName) {
        if ($IsYaml) {
            # 从YAML文件获取环境名称
            $yamlContent = Get-Content $FilePath | Select-String "name: (.+)"
            if ($yamlContent -and $yamlContent.Matches.Groups.Count -gt 1) {
                $defaultName = $yamlContent.Matches.Groups[1].Value
                $askForName = Read-Host "Do you want to use the default name '$defaultName' from YAML? (y/n)"
                if ($askForName.ToLower() -eq "y") {
                    $CustomEnvName = $defaultName
                }
            }
        }
        
        if (-not $CustomEnvName) {
            $CustomEnvName = Read-Host "Please enter a name for the new environment"
            if (-not $CustomEnvName) {
                $CustomEnvName = "py_env_$(Get-Date -Format 'yyyyMMdd_HHmmss')"
                Write-Host "Using auto-generated name: $CustomEnvName" -ForegroundColor Yellow
            }
        }
    }
    
    # 从YAML文件获取Python版本或让用户指定
    if (-not $PythonVersion) {
        if ($IsYaml) {
            # 尝试从YAML获取Python版本
            $content = Get-Content $FilePath -Raw
            if ($content -match "python=(\d+\.\d+)") {
                $detectedVersion = $matches[1]
                Write-Host "Detected Python version $detectedVersion in YAML file" -ForegroundColor Cyan
                $PythonVersion = $detectedVersion
            }
        }
        
        if (-not $PythonVersion) {
            $PythonVersion = Read-Host "Please enter Python version (e.g. 3.9)"
            if (-not $PythonVersion) {
                $PythonVersion = "3.9"
                Write-Host "Using default Python version: $PythonVersion" -ForegroundColor Yellow
            }
        }
    }
    
    # 确保Python版本格式正确
    if (-not ($PythonVersion -match "^python=")) {
        if ($PythonVersion -match "^\d+\.\d+$") {
            $PythonVersion = "python=$PythonVersion"
        } else {
            $PythonVersion = "python=$PythonVersion"
        }
    }
    
    Write-Host "Creating environment '$CustomEnvName' with $PythonVersion..." -ForegroundColor Cyan
    
    # Configure SSL (配置SSL)
    Configure-SSL
    
    # Configure channels (配置源)
    Write-Host "Configuring channels..." -ForegroundColor Cyan
    Configure-Channels -UseTsinghua $UseTsinghua
    
    # 创建包含Python的基础环境
    $createCmd = "conda create -n `"$CustomEnvName`" $PythonVersion --yes"
    Write-Host "Executing: $createCmd" -ForegroundColor Gray
    Invoke-Expression $createCmd
    
    if ($LASTEXITCODE -ne 0) {
        Write-Host "Failed to create base environment" -ForegroundColor Red
        return $false
    }
    
    # 根据文件类型决定如何安装包
    if ($IsYaml) {
        # 安装YAML文件中除Python外的其他包
        Write-Host "Installing packages from YAML..." -ForegroundColor Yellow
        
        # 从YAML文件解析依赖项
        $content = Get-Content $FilePath -Raw
        if ($content -match "dependencies:\s*\n((?:\s*-\s*.+\n)+)") {
            $dependencies = $matches[1] -split "`n" | Where-Object { $_ -match "^\s*-\s*(.+)$" } | ForEach-Object { $matches[1].Trim() }
            
            # 逐个安装包
            foreach ($dep in $dependencies) {
                if ($dep -notmatch "^python=?") {  # 跳过Python
                    $installCmd = "conda install -n `"$CustomEnvName`" $dep --yes --no-deps"
                    Write-Host "Installing $dep..." -ForegroundColor Cyan
                    Write-Host "Executing: $installCmd" -ForegroundColor Gray
                    
                    Invoke-Expression $installCmd
                    if ($LASTEXITCODE -ne 0) {
                        Write-Host "Warning: Failed to install $dep, skipping..." -ForegroundColor Yellow
                    }
                    else {
                        Write-Host "Successfully installed $dep" -ForegroundColor Green
                    }
                }
            }
        }
    }
    else {
        # 使用pip安装requirements.txt中的包
        Write-Host "Installing packages from requirements.txt..." -ForegroundColor Yellow
        
        # 使用cmd运行conda activate和pip install命令
        $pipCmd = "cmd.exe /c `"conda activate $CustomEnvName && pip install -r `"$FilePath`"`""
        Write-Host "Executing: conda activate $CustomEnvName && pip install -r `"$FilePath`"" -ForegroundColor Gray
        Invoke-Expression $pipCmd
        
        if ($LASTEXITCODE -ne 0) {
            Write-Host "Warning: Some packages may have failed to install" -ForegroundColor Yellow
        } else {
            Write-Host "Successfully installed packages from requirements.txt" -ForegroundColor Green
        }
    }
    
    Write-Host "Environment creation completed" -ForegroundColor Green
    return $true
}

# Main function (主函数)
function Main {
    # Export environment if requested (如果请求导出环境)
    if ($Export) {
        if (-not $EnvName) {
            $envs = List-CondaEnvironments
            
            Write-Host "Available conda environments:" -ForegroundColor Cyan
            for ($i = 0; $i -lt $envs.Count; $i++) {
                Write-Host "$($i+1). $($envs[$i].Name)" -ForegroundColor White
            }
            
            try {
                $choice = [int](Read-Host "Enter the number of the environment to export")
                if ($choice -ge 1 -and $choice -le $envs.Count) {
                    $EnvName = $envs[$choice-1].Name
                }
                else {
                    Write-Host "Invalid selection" -ForegroundColor Red
                    return
                }
            }
            catch {
                Write-Host "Please enter a valid number" -ForegroundColor Red
                return
            }
        }
        
        $exportType = Read-Host "Export as (1) requirements.txt, (2) environment.yml, or (3) both? [1-3]"
        
        switch ($exportType) {
            "1" { Export-Environment -EnvName $EnvName }
            "2" { Export-EnvironmentYAML -EnvName $EnvName }
            "3" { 
                Export-Environment -EnvName $EnvName
                Export-EnvironmentYAML -EnvName $EnvName 
            }
            default { 
                Write-Host "Invalid choice, exporting both formats" -ForegroundColor Yellow
                Export-Environment -EnvName $EnvName
                Export-EnvironmentYAML -EnvName $EnvName 
            }
        }
        
        return
    }
    
    # Export all environments if requested (如果请求导出所有环境)
    if ($ExportAll) {
        Export-AllEnvironments
        return
    }
    
    # Handle list parameter (处理列表参数)
    if ($List) {
        $available_envs = List-AvailableEnvironments
        if ($available_envs.Count -eq 0) {
            return
        }
        
        Write-Host "Available environment files:" -ForegroundColor Cyan
        foreach ($env in $available_envs) {
            Write-Host "- $env" -ForegroundColor White
        }
        return
    }
    
    # 处理requirements.txt创建
    if ($RequirementsFile) {
        $requirementsPath = $RequirementsFile
        
        # 如果路径不是绝对路径，就当作相对路径处理
        if (-not ([System.IO.Path]::IsPathRooted($requirementsPath))) {
            # 尝试在当前目录查找
            if (-not (Test-Path $requirementsPath)) {
                # 尝试在conda_environments目录查找
                $envDirPath = Join-Path -Path (Get-Location) -ChildPath $ENV_DIR
                $potentialPath = Join-Path -Path $envDirPath -ChildPath $requirementsPath
                
                if (Test-Path $potentialPath) {
                    $requirementsPath = $potentialPath
                }
            }
        }
        
        Create-Environment -FilePath $requirementsPath -CustomEnvName $EnvName -PythonVersion $PythonVersion -UseTsinghua (-not $NoTsinghua) -IsYaml $false
        return
    }
    
    # 处理YAML文件创建
    if ($EnvFile) {
        # 如果路径不是绝对路径，就当作相对路径处理
        $yamlPath = $EnvFile
        if (-not ([System.IO.Path]::IsPathRooted($yamlPath))) {
            # 尝试在当前目录查找
            if (-not (Test-Path $yamlPath)) {
                # 尝试在conda_environments目录查找
                $envDirPath = Join-Path -Path (Get-Location) -ChildPath $ENV_DIR
                $potentialPath = Join-Path -Path $envDirPath -ChildPath $yamlPath
                
                if (Test-Path $potentialPath) {
                    $yamlPath = $potentialPath
                }
            }
        }
        
        Create-Environment -FilePath $yamlPath -CustomEnvName $EnvName -PythonVersion $PythonVersion -UseTsinghua (-not $NoTsinghua) -IsYaml $true
        return
    }
    
    # 如果没有指定环境文件，显示交互式选择
    if (-not $EnvFile -and -not $RequirementsFile) {
        $available_envs = List-AvailableEnvironments
        if ($available_envs.Count -eq 0) {
            return
        }
        
        Write-Host "Please select an environment file:" -ForegroundColor Cyan
        for ($i = 0; $i -lt $available_envs.Count; $i++) {
            Write-Host "$($i+1). $($available_envs[$i])" -ForegroundColor White
        }
        
        try {
            $choice = [int](Read-Host "Enter the number")
            if ($choice -ge 1 -and $choice -le $available_envs.Count) {
                $EnvFile = $available_envs[$choice-1]
                
                # 获取YAML文件完整路径
                $yamlPath = Join-Path -Path $ENV_DIR -ChildPath $EnvFile
                
                # 创建环境
                Create-Environment -FilePath $yamlPath -CustomEnvName $EnvName -PythonVersion $PythonVersion -UseTsinghua (-not $NoTsinghua) -IsYaml $true
            }
            else {
                Write-Host "Invalid selection" -ForegroundColor Red
                return
            }
        }
        catch {
            Write-Host "Please enter a valid number" -ForegroundColor Red
            return
        }
    }
}

# Execute main function (执行主函数)
Main 